import torch
from scipy.optimize import linear_sum_assignment
from torch import nn

from util.box_ops import box_cxcywh_to_xyxy, generalized_box_iou
from .matcherhoi import HungarianMatcherHOI as matcher

def build_matcher(args):
    return HungarianMatcherHOI(cost_obj_class=args.set_cost_obj_class, cost_verb_class=args.set_cost_verb_class,
                               cost_bbox=args.set_cost_bbox, cost_giou=args.set_cost_giou,
                               cost_hoi_class=args.set_cost_hoi)
    
class HungarianMatcherHOI(matcher):

    def __init__(self, cost_obj_class: float = 1, cost_verb_class: float = 1, cost_bbox: float = 1,
                 cost_giou: float = 1, cost_hoi_class: float = 1):
        super(HungarianMatcherHOI, self).__init__(cost_obj_class, cost_verb_class, cost_bbox, cost_giou)

        
        self.cost_hoi_class = cost_hoi_class
        assert cost_obj_class != 0 or cost_verb_class != 0 or cost_bbox != 0 or cost_giou != 0, 'all costs cant be 0'

    @torch.no_grad()
    def forward(self, outputs, targets):
        C = self.get_C(outputs, targets)       
        sizes = [len(v['obj_labels']) for v in targets]
        #   list长度为N，i是pred，j是GT,    寻找loss最小的(不会重复)
        indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]
        indices = [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
        return indices


    @torch.no_grad()
    def get_C(self, outputs, targets):
        bs, num_queries = outputs['pred_obj_logits'].shape[:2]         
        C = self.get_cost_obj_class(outputs['pred_obj_logits'], targets, 
                                    self.obj_alpha, self.obj_gamma, self.neg_inds_beta, self.hard_neg_inds_beta) + \
            self.get_cost_bbox(outputs['pred_sub_boxes'], outputs['pred_obj_boxes'], targets) + \
            self.get_cost_giou(outputs['pred_sub_boxes'], outputs['pred_obj_boxes'], targets)     
        if 'pred_hoi_logits' in outputs.keys():
            C += self.get_cost_hoi_class(outputs['pred_hoi_logits'], targets, 
                                     self.verb_alpha, self.verb_gamma, self.neg_inds_beta, self.hard_neg_inds_beta)
        else:
            C +=  self.get_cost_verb_class(outputs['pred_verb_logits'], targets, 
                                     self.verb_alpha, self.verb_gamma, self.neg_inds_beta, self.hard_neg_inds_beta)
        C = C.view(bs, num_queries, -1).cpu()
        return C    



    def get_cost_hoi_class(self, pred_hoi_logits, targets, alpha=0.5, gamma=2, neg_inds_beta=0, hard_neg_inds_beta=3):
        out_hoi_prob = pred_hoi_logits.flatten(0, 1).sigmoid()
        tgt_hoi_labels = torch.cat([v['hoi_labels'] for v in targets])
        tgt_hoi_labels = tgt_hoi_labels.permute(1, 0)
        
        if self.loss_type=='ce_loss':  
            pos = out_hoi_prob.matmul(tgt_hoi_labels)
            pos_sum = tgt_hoi_labels.sum(dim=0, keepdim=True) + 1e-4
            neg = (1 - out_hoi_prob).matmul(1 - tgt_hoi_labels)
            neg_num = ((1 - tgt_hoi_labels).sum(dim=0, keepdim=True) + 1e-4)
            
            cost_hoi_class = -(pos / pos_sum + neg / neg_num) / 2
            
        elif self.loss_type=='focal_loss':
            #out_hoi_prob = torch.clamp(out_hoi_prob, 1e-12) 
            pos_pred = alpha * ((1 - out_hoi_prob) ** gamma) * (-(out_hoi_prob + 1e-8).log())
            pos_inds = tgt_hoi_labels.eq(1).float()
            pos_target = pos_inds 
            pos_hoi_cost_class = pos_pred.matmul(pos_target)
            pos_tgt_hoi_labels_sum = pos_inds.sum(dim=0, keepdim=True) + 1e-4
            
            neg_pred= (1 - alpha) * (out_hoi_prob ** gamma) * (-(1 - out_hoi_prob + 1e-8).log())
            neg_inds = (tgt_hoi_labels.lt(1) & tgt_hoi_labels.gt(-1)).float()
            hard_neg_inds = (tgt_hoi_labels.eq(-1)).float()
            neg_weights = torch.pow(1 - tgt_hoi_labels, neg_inds_beta) * neg_inds + \
                        torch.pow(1 - tgt_hoi_labels, hard_neg_inds_beta) * hard_neg_inds #   2
            neg_target = neg_weights 
            neg_hoi_cost_class = neg_pred.matmul(neg_target)
            neg_tgt_hoi_labels_sum = (neg_inds + hard_neg_inds).sum(dim=0, keepdim=True) + 1e-4
            cost_hoi_class = (pos_hoi_cost_class  + neg_hoi_cost_class) / pos_tgt_hoi_labels_sum
                    
        #   100N, tgtN
        C = self.cost_hoi_class * cost_hoi_class
        return C